# ------------------------------------------------------------------------
# Deformable DETR
# Copyright (c) 2020 SenseTime. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
# Modified from DETR (https://github.com/facebookresearch/detr)
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
# ------------------------------------------------------------------------

"""
Deformable DETR model and criterion classes.
"""
import torch
from torch import nn
import math

from util.misc import (NestedTensor, nested_tensor_from_tensor_list)
from .module.modules import MLP

from .backbones.backbone import build_backbone
from .necks.smca_transformer import build_transformer
from .matcher.matcher import build_matcher
from .losses.setcriterion import SetCriterion
from .postprocess.smca_postprocess import PostProcess
import numpy as np

def build(args):
    device = torch.device(args.device)

    num_classes = 20 if args.dataset_file != 'coco' else 91
    if args.dataset_file == "coco_panoptic":
        num_classes = 250

    backbone = build_backbone(args)

    transformer = build_transformer(args)

    model = SMCA(
        backbone,
        transformer,
        num_classes=num_classes,
        num_queries=args.num_queries,
        aux_loss=args.aux_loss,
    )
    # if args.masks:
    #     model = DETRsegm(model, freeze_detr=(args.frozen_weights is not None))
    matcher = build_matcher(args)
    weight_dict = {'loss_ce': 2, 'loss_bbox': args.bbox_loss_coef}
    weight_dict['loss_giou'] = args.giou_loss_coef
    if args.masks:
        weight_dict["loss_mask"] = args.mask_loss_coef
        weight_dict["loss_dice"] = args.dice_loss_coef
    # TODO this is a hack
    if args.aux_loss:
        aux_weight_dict = {}
        for i in range(args.dec_layers - 1):
            aux_weight_dict.update({k + f'_{i}': v for k, v in weight_dict.items()})
        weight_dict.update(aux_weight_dict)

    losses = ['labels', 'boxes', 'cardinality']
    # if args.masks:
    #     losses += ["masks"]
    criterion = SetCriterion(num_classes, matcher=matcher, weight_dict=weight_dict,
                             eos_coef=args.eos_coef, losses=losses)
    criterion.to(device)
    postprocessors = {'bbox': PostProcess()}
    # if args.masks:
    #     postprocessors['segm'] = PostProcessSegm()
    #     if args.dataset_file == "coco_panoptic":
    #         is_thing_map = {i: i <= 90 for i in range(201)}
    #         postprocessors["panoptic"] = PostProcessPanoptic(is_thing_map, threshold=0.85)

    return model, criterion, postprocessors


class SMCA(nn.Module):
    """ This is the DETR module that performs object detection """
    def __init__(self, backbone, transformer, num_classes, num_queries, aux_loss=False, num_verb_classes=0):
        super().__init__()
        
        self.num_queries = num_queries
        self.transformer = transformer
        hidden_dim = transformer.d_model
        self.hidden_dim = hidden_dim
        self.query_embed = nn.Embedding(num_queries, hidden_dim)
        self.input_proj = nn.Conv2d(backbone.num_channels[-1], hidden_dim, kernel_size=1)
        self.backbone = backbone
        self.aux_loss = aux_loss
        self.build_head(num_classes, num_verb_classes)

 
    def build_head(self, num_classes, num_verb_classes):
        hidden_dim = self.hidden_dim
        self.class_embed = nn.Linear(hidden_dim, num_classes + 1)
        self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
        
        prior_prob = 0.01
        bias_value = -math.log((1 - prior_prob) / prior_prob)
        self.class_embed.bias.data = torch.ones(num_classes) * bias_value
        nn.init.constant_(self.bbox_embed.layers[-1].weight.data, 0)
        nn.init.constant_(self.bbox_embed.layers[-1].bias.data, 0)


    def forward(self, samples: NestedTensor, targets=None):
        srcs, masks, poses, h_w = self.backbone_forward(samples, targets)
        outputs = self.neck_forward(srcs, masks, poses, self.query_embed.weight, h_w) 
        out = self.head_forward(outputs)
        out.update(outputs)
        return out

    def backbone_forward(self, samples: NestedTensor, targets=None):
        if isinstance(samples, (list, torch.Tensor)):
            samples = nested_tensor_from_tensor_list(samples)
        features, pos = self.backbone(samples)

        h_w = torch.stack([torch.stack([inst['size'] for inst in targets])[:, 1],
                           torch.stack([inst['size'] for inst in targets])[:, 0]], dim=-1)
        h_w = h_w.unsqueeze(0)

        srcs = []
        masks = []
        poses = []
        for j, feat in enumerate(features):
            if j== len(features)-1:
                src, mask = feat.decompose()
                srcs.append(self.input_proj(src))
                masks.append(mask)
                poses.append(pos[j])
                assert mask is not None

        return srcs, masks, poses, h_w
    
    def neck_forward(self, srcs, masks, poses, query_embed, h_w):
        ###neck###
        #   src:               N, 2048, 35, 35 -> N, 256, 35, 35
        #   mask:              N, 35, 35 
        #   query_embed:       100, N, 256
        #   pos:               256, 35, 35
        #-> hs                 6, N, 100, 256
        #{'src':src,'out_query':hs,'memory':memory}
        outputs = self.transformer(srcs, masks, query_embed, poses, h_w)
        return outputs

    def head_forward(self, outputs):
        #   hs                 6, N, 100, 256
        #-> outputs_obj_class  N, 100, 80
        x = outputs['out_query']
        points = outputs['points']

        num_decoder = x.shape[0]

        outputs_class = self.class_embed(x)
        outputs_coord = self.bbox_embed(x)

        points = points.unsqueeze(0).repeat(num_decoder, 1, 1, 1)

        outputs_coord[..., :2] = outputs_coord[..., :2] + points
        outputs_coord = outputs_coord.sigmoid()
        out = {'pred_logits': outputs_class[-1], 
               'pred_boxes': outputs_coord[-1]}
        if self.aux_loss and self.training:
            out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord)
        return out

    @torch.jit.unused
    def _set_aux_loss(self, outputs_class, outputs_coord):
        # this is a workaround to make torchscript happy, as torchscript
        # doesn't support dictionary with non-homogeneous values, such
        # as a dict having both a Tensor and a list.
        return [{'pred_logits': a, 'pred_boxes': b}
                for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]



def bias_init_with_prob(prior_prob):
    """initialize conv/fc bias value according to giving probablity."""
    bias_init = float(-np.log((1 - prior_prob) / prior_prob))
    return bias_init
